import io
import os
import sys
import threading
import time
import traceback
from itertools import count
from typing import Any, BinaryIO, Callable, ContextManager, Optional, TextIO, Union

from paramiko import SFTPClient, SSHClient, WarningPolicy
from paramiko.ssh_exception import NoValidConnectionsError

from .config import SSHConfig


def ssh_connect(
    ssh_config: SSHConfig, *,
    initial_delay: float = 3,
    connect_timeout: float = 3,
    retry_interval: float = 3,
    retry_times: int = 20,
    verbose: bool = True
) -> SSHClient:
    if verbose:
        print(f"Connecting to {ssh_config.url}", file=sys.stderr)

    ssh_client = SSHClient()
    ssh_client.set_missing_host_key_policy(WarningPolicy)

    time.sleep(initial_delay)
    for i in count():
        try:
            ssh_client.connect(
                hostname=ssh_config.hostname,
                port=ssh_config.port,
                username=ssh_config.username,
                pkey=ssh_config.private_key,
                key_filename=ssh_config.private_key_filepath,
                timeout=connect_timeout)
            return ssh_client
        except (TimeoutError, NoValidConnectionsError):
            if i >= retry_times:
                raise
            traceback.print_exc(file=sys.stderr)
            time.sleep(retry_interval)
            if verbose:
                print("Retrying...", file=sys.stderr)
            continue


# execute utils

def ssh_call(
    ssh_client: SSHClient,
    command: str, *,
    encoding: Optional[str] = 'utf8',
    raise_on_error: bool = True,
) -> Union[str, bytes]:
    stdin, stdout, stderr = ssh_client.exec_command(command)
    result = stdout.read()

    code = stdout.channel.recv_exit_status()
    if raise_on_error and code != 0:
        raise RuntimeError(f"Command exited with non-zero code: {code}")

    if encoding is not None:
        result = result.decode(encoding)
    return result


def ssh_execute(ssh_client: SSHClient, command: str, *, raise_on_error: bool = True) -> int:
    stdin, stdout, stderr = ssh_client.exec_command(command)
    attach_outputs(stdout, stderr)

    code = stdout.channel.recv_exit_status()
    if raise_on_error and code != 0:
        raise RuntimeError(f"Command exited with non-zero code: {code}")
    return code


def attach_outputs(stdout: BinaryIO | Any | None = None, stderr: BinaryIO | Any | None = None):
    if stdout is not None:
        stdout_thread_func = lambda: io_transfer(stdout, sys.stdout.buffer, by_line=True)
        stdout_thread = threading.Thread(target=stdout_thread_func, daemon=True)
        stdout_thread.start()
    else:
        stdout_thread = None

    if stderr is not None:
        stderr_thread_func = lambda: io_transfer(stderr, sys.stderr.buffer, by_line=True)
        stderr_thread = threading.Thread(target=stderr_thread_func, daemon=True)
        stderr_thread.start()
    else:
        stderr_thread = None

    if stdout_thread is not None:
        while stdout_thread.is_alive():
            stdout_thread.join(1)
    if stderr_thread is not None:
        while stderr_thread.is_alive():
            stderr_thread.join(1)


# upload utils

UploadSource = Union[bytes, BinaryIO, TextIO, str, os.PathLike, Callable, ContextManager]


def ssh_upload(ssh_client: SSHClient, source: UploadSource, dst_path: str, *, verbose: bool = True):
    with ssh_client.open_sftp() as sftp:
        sftp_upload(sftp, source, dst_path, verbose=verbose)


def sftp_upload(sftp_client: SFTPClient, source: UploadSource, dst_path: str, *, verbose: bool = True):
    if isinstance(source, (str, os.PathLike)):
        if os.path.isfile(source):
            sftp_mkdirs(sftp_client, os.path.dirname(dst_path), verbose=verbose)
            sftp_client.put(source, dst_path)
            if verbose:
                print(f"Uploaded {source} -> {dst_path}", file=sys.stderr)
        elif os.path.isdir(source):
            sftp_mkdirs(sftp_client, dst_path, verbose=verbose)
            for sub_path in os.listdir(source):
                sub_src_path = os.path.join(source, sub_path)
                sub_dst_path = dst_path + '/' + sub_path
                sftp_upload(sftp_client, sub_src_path, sub_dst_path, verbose=verbose)
        else:
            raise FileNotFoundError(f"Source path not found: {source}")
        return

    if isinstance(source, TextIO):
        sftp_upload(sftp_client, source.buffer, dst_path, verbose=False)
        return

    if isinstance(source, bytes):
        sftp_upload(sftp_client, io.BytesIO(source), dst_path, verbose=False)
        if verbose:
            print(f"Uploaded <content> -> {dst_path}", file=sys.stderr)
        return

    if isinstance(source, (BinaryIO, io.BytesIO)):
        sftp_mkdirs(sftp_client, os.path.dirname(dst_path), verbose=verbose)
        sftp_client.putfo(source, dst_path)
        if verbose:
            print(f"Uploaded <stream> -> {dst_path}", file=sys.stderr)
        return

    if isinstance(source, ContextManager):
        with source as tmp_source:
            sftp_upload(sftp_client, tmp_source, dst_path)
        return

    if isinstance(source, Callable):
        tmp_source = source()
        sftp_upload(sftp_client, tmp_source, dst_path)
        return

    raise TypeError(f"Unexpected type of source: {type(source)}")


def sftp_mkdirs(sftp_client: SFTPClient, dst_path: str, *, verbose: bool = True):
    if sftp_exists(sftp_client, dst_path):
        return

    parent_path = os.path.dirname(dst_path)
    sftp_mkdirs(sftp_client, parent_path, verbose=verbose)

    sftp_client.mkdir(dst_path)
    if verbose:
        print(f"Created directory: {dst_path}", file=sys.stderr)


def sftp_exists(sftp_client: SFTPClient, dst_path: str):
    try:
        sftp_client.stat(dst_path)
        return True
    except FileNotFoundError:
        return False


# download utils

def ssh_download(ssh_client: SSHClient, src_path: str, dst_path: str, *, verbose: bool = True):
    with ssh_client.open_sftp() as sftp:
        sftp_download(sftp, src_path, dst_path, verbose=verbose)


def sftp_download(sftp_client: SFTPClient, src_path: str, dst_path: str, *, verbose: bool = True):
    with sftp_client.open(src_path, 'rb') as src, open(dst_path, 'wb') as dst:
        # noinspection PyTypeChecker
        io_transfer(src, dst)
    if verbose:
        print(f"\tDownloaded {src_path} -> {dst_path}", file=sys.stderr)


# io utils

def io_transfer(
    in_io: BinaryIO,
    out_io: BinaryIO, *,
    by_line: bool = False,
    buffer_size: int = 1024,
):
    if by_line:
        in_io = io.TextIOWrapper(in_io, newline='')
    while True:
        if by_line:
            piece = in_io.readline(buffer_size)
        else:
            piece = in_io.read(buffer_size)
        if not piece:
            break
        if isinstance(piece, str):
            piece = piece.encode()
        out_io.write(piece)
